# -*- Coding: UTF-8 -*-
# main.py
# @author: SongLiangCao
# @email: 2023733547@qq.com
# @description: none
# @created: 2021-12-16T19:25:17.130Z+08:00
# @last-modified: 2022-01-10T22:32:21.122Z+08:00
#

import time
import numpy as np
import matplotlib.pyplot as plt
from ConvNet import ConvNet, Train
from mnist import load_mnist


def main():
    start = time.clock()
    (train_images, train_label), (test_images, test_label) = load_mnist()
    convNet = ConvNet()
    train_images = train_images.reshape(train_images.shape[0], 1, 28, 28)
    test_images = test_images.reshape(test_images.shape[0], 1, 28, 28)
    print("train begin")
    print("=====================")
    trainer = Train(convNet, train_images, train_label, test_images, test_label, epochs=15)
    trainer.train()
    print("end")

    # plot the data
    x = np.arange(len(trainer.train_acc_list))
    plt.plot(x, trainer.train_acc_list, label='train', markevery=10)
    plt.plot(x, (trainer.test_acc_list), label='test', markevery=10)
    plt.xlabel("epochs")
    plt.ylabel("accuracy")
    plt.ylim(0, 1.0)
    plt.legend(loc='lower right')
    plt.show()
    x = np.arange(len(trainer.train_loss_list))
    plt.plot(x, (trainer.train_loss_list), label='loss', markevery=400)
    plt.xlabel("epochs")
    plt.ylabel("loss")
    plt.show()
    end = time.clock()
    print('Running time: %s Seconds' % (end-start))


if __name__ == '__main__':
    main()
